#!python3

import sys
import hashlib
from base64 import b64encode, b64decode
import math

prooffile = "proof.txt"

MAXHEIGHT = 20


class MerkleProof:
    def __init__(self, leaf, pos, path):
        self.leaf = leaf
        self.pos  = pos
        self.path = path


def hash_leaf(leaf):
    sha256 = hashlib.sha256()
    sha256.update(b"leaf:")
    sha256.update(leaf)
    return sha256.digest()


def hash_internal_node(left, right):
    sha256 = hashlib.sha256()
    sha256.update(b"node:")
    sha256.update(left)
    sha256.update(right)
    return sha256.digest()


def gen_merkle_proof(leaves, pos):
    height = math.ceil(math.log(len(leaves),2))
    assert height < MAXHEIGHT, "Too many leaves."

    state = list(map(hash_leaf, leaves))  

    padlen = (2**height)-len(leaves)
    state += [b"\x00"] * padlen

    path = []

    level_pos = pos

    for level in range(height):
        if level_pos % 2 == 0:
            sibling_pos = level_pos + 1
        else:
            sibling_pos = level_pos - 1
        
        path.append(state[sibling_pos])
        
        parent_pos = level_pos // 2
        
        new_state = []
        for i in range(0, len(state), 2):
            new_state.append(hash_internal_node(state[i], state[i+1]))
        
        state = new_state
        level_pos = parent_pos

    return path


def write_proof(filename, proof:MerkleProof):
    fp = open(filename, "w")
    print("leaf position: {pos:d}".format(pos=proof.pos), file=fp)
    print("leaf value: \"{leaf:s}\"".format(leaf=proof.leaf.decode('utf-8')), file=fp)
    print("Hash values in proof:", file=fp)
    for i in range(len(proof.path)):
        print("  {:s}".format(b64encode(proof.path[i]).decode('utf-8')), file=fp)
    fp.close()


if __name__ == "__main__":
    leaves = [b"data item " + str(i).encode() for i in range(1000)]
    print('\nI generated 1000 leaves for a Merkle tree of height 10.')

    pos = 743
    path  = gen_merkle_proof(leaves, pos)
    proof = MerkleProof(leaves[pos], pos, path)

    write_proof(prooffile, proof)

    print('I generated a Merkle proof for leaf #{} in file {}\n'.format(pos,prooffile))
    sys.exit(0)



